# 导入必要的库
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error


# 数据预处理函数
def getTrainSetAndTestSet(DataPath):
    data = pd.read_csv(DataPath)
    X = data[['AT', 'V', 'AP', 'RH']]  # 特征
    y = data['PE']  # 标签（注意改为Series格式）
    # 划分训练集和测试集（3:1比例）
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=1)
    # 输出数据集维度
    print("训练集维度：X", X_train.shape, "y", y_train.shape)
    print("测试集维度：X", X_test.shape, "y", y_test.shape)
    return X_train, X_test, y_train, y_test


# 训练线性回归模型
def TrainLinearRegression(X_train, y_train):
    linreg = LinearRegression()
    linreg.fit(X_train, y_train)
    # 输出模型参数
    print("回归系数：", linreg.coef_)
    print("截距项：", linreg.intercept_)
    return linreg


# 模型评估
def EvaluationModel(linreg, X_test, y_test):
    y_pred = linreg.predict(X_test)
    # 计算MSE和RMSE
    mse = mean_squared_error(y_test, y_pred)
    rmse = mse ** 0.5
    print("均方误差(MSE):", mse)
    print("均方根误差(RMSE):", rmse)
    return y_pred


# 可视化结果
def Visualization(y_test, y_pred):
    plt.figure(figsize=(8, 6))
    plt.scatter(y_test, y_pred, alpha=0.5)
    plt.plot([y_test.min(), y_test.max()],
             [y_test.min(), y_test.max()],
             'k--', lw=2)  # 黑色虚线参考线
    plt.xlabel('Measured')
    plt.ylabel('Predicted')
    plt.title('Actual vs Predicted')
    plt.show()


# 主程序
if __name__ == "__main__":
    # 数据路径需根据实际文件位置修改
    DataPath = 'Folds5x2_pp.csv'

    # 数据预处理
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(DataPath)

    # 训练模型
    model = TrainLinearRegression(X_train, y_train)

    # 模型评估
    y_pred = EvaluationModel(model, X_test, y_test)

    # 可视化
    Visualization(y_test, y_pred)